#!/usr/bin/env python
"""$Id$"""

import os
import math
from math import sin, cos, radians, degrees, atan2, sqrt, exp
import ctypes
from pyglet.gl import *
from pyglet import window
from pyglet.window import key
import pyglet.clock
import ode
import terrain
from skybox import SkyBox
import data


AIR_DENSITY = 1.2 # kg/m^3 (70 degrees at std atmospheric pressure)
AIR_TEMP = 288.15 # K
PAUSED = False


def resize(width, height):
	if height==0:
		height=1
	glViewport(0, 0, width, height)
	glMatrixMode(GL_PROJECTION)
	glLoadIdentity()
	gluPerspective(60, 1.0*width/height, .1, 30000.0)
	glMatrixMode(GL_MODELVIEW)
	glLoadIdentity()

def init():
	glClearColor(0.7, 0.7, 1.0, 1.0)
	glClearDepth(1.0)
	glClearStencil(0)	
	glEnable(GL_COLOR_MATERIAL)
	glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE)
	glEnable(GL_LIGHTING)
	glEnable(GL_LIGHT0)
	#glLightModeli(GL_LIGHT_MODEL_TWO_SIDE, GL_TRUE)
	fourfv = ctypes.c_float * 4
	#glLightfv(GL_LIGHT0, GL_POSITION, fourfv(100, 200, 200, 0))
	glLightfv(GL_LIGHT0, GL_POSITION, fourfv(0, 1, 0, 0))
	glLightfv(GL_LIGHT0, GL_AMBIENT, fourfv(0.2, 0.2, 0.2, 1.0))
	glLightfv(GL_LIGHT0, GL_DIFFUSE, fourfv(0.7, 0.7, 0.7, 1.0))
	glLightfv(GL_LIGHT0, GL_SPECULAR, fourfv(0.05, 0.05, 0.05, 1.0))
	glEnable(GL_DEPTH_TEST)
	glDepthFunc(GL_LEQUAL)
	glEnable(GL_CULL_FACE)
	glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST)
	glEnable(GL_BLEND)
	glShadeModel(GL_SMOOTH)
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
	#glEnable(GL_MULTISAMPLE)

def draw_ground():
	'''
	glBegin(GL_QUADS)				
	glColor3f(1.0, 0.95, 0.95)
	glVertex4f(-1000, 0, -1000, 1)
	glVertex4f(-1000, 0, 1000, 1)
	glVertex4f(1000, 0, 1000, 1)
	glVertex4f(1000, 0, -1000, 1)
	glEnd()
	'''
	glBegin(GL_LINES)
	glColor3f(0, 0, 0)
	for i in range(1000):
		glVertex3f(-10 * i, 0, -10000)
		glVertex3f(-10 * i, 0, 10000)
		glVertex3f(10 * i, 0, -10000)
		glVertex3f(10 * i, 0, 10000)
		glVertex3f(-10000, 0, -10 * i)
		glVertex3f(10000, 0, -10 * i)
		glVertex3f(-10000, 0, 10 * i)
		glVertex3f(10000, 0, 10 * i)
	glEnd()

circle_pts = [(a, math.cos(a*math.pi*2.0/40.0), math.sin(a*math.pi*2.0/40.0)) 
	for a in range(41)]

def draw_rotor(radius, cone=0.0, rot=[0]):
	glDisable(GL_CULL_FACE)
	glBegin(GL_TRIANGLE_FAN)
	glVertex3f(0.0, 0.0, 0.0)
	for r, L1, L2 in circle_pts:
		if r % 10 == rot[0]:
			glColor4f(0.0, 0.0, 0.0, 0.3)
		else:
			glColor4f(0.0, 0.0, 0.0, 0.005)
		glVertex3f(L1 * radius, cone, L2 * radius)
	rot[0] = (rot[0] - 1) % 10
	glEnd()
	glEnable(GL_CULL_FACE)
	
def draw():
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
	glLoadIdentity()			
	
	glTranslatef(0.0,-2.0,-5.0)

	#glRotatef(90,1.0,0.0,0.0)			
	glRotatef(rtri,0.0,1.0,0.0)			

	glBegin(GL_QUADS)				

	glColor4f(0.0, 0.0, 0.0, 1.0) # blade
	glVertex3f(0.0, 1.0, 0.0)		
	glVertex3f(-1.0, 1.0, 0.0)
	glColor4f(0.0, 0.0, 0.0, 0.5)	
	glVertex3f(-1.0, 1.0, 0.07)
	glVertex3f(-0.2, 1.0, 0.07)
	glColor4f(1.0, 1.0, 0.0, 1.0) # tip	
	glVertex3f(-1.0, 1.0, 0.0)
	glVertex3f(-1.04, 1.0, 0.0)
	glColor4f(1.0, 1.0, 0.0, 0.0)	
	glVertex3f(-1.04, 1.0, 0.07)
	glVertex3f(-1.0, 1.0, 0.07)

	glColor4f(0.0, 0.0, 0.0, 1.0) # blade
	glVertex3f(0.0, 1.0, 0.0)		
	glVertex3f(0.0, 1.0, 1.0)
	glColor4f(0.0, 0.0, 0.0, 0.5)	
	glVertex3f(0.07, 1.0, 1.0)
	glVertex3f(0.07, 1.0, 0.2)
	glColor4f(1.0, 1.0, 0.0, 1.0) # tip	
	glVertex3f(0.0, 1.0, 1.0)
	glVertex3f(0.0, 1.0, 1.04)
	glColor4f(1.0, 1.0, 0.0, 0.0)	
	glVertex3f(0.07, 1.0, 1.04)
	glVertex3f(0.07, 1.0, 1.0)

	glColor4f(0.0, 0.0, 0.0, 1.0) # blade
	glVertex3f(0.0, 1.0, 0.0)		
	glVertex3f(1.0, 1.0, 0.0)
	glColor4f(0.0, 0.0, 0.0, 0.5)	
	glVertex3f(1.0, 1.0, -0.07)
	glVertex3f(0.2, 1.0, -0.07)
	glColor4f(1.0, 1.0, 0.0, 1.0) # tip	
	glVertex3f(1.0, 1.0, 0.0)
	glVertex3f(1.04, 1.0, 0.0)
	glColor4f(1.0, 1.0, 0.0, 0.0)	
	glVertex3f(1.04, 1.0, -0.07)
	glVertex3f(1.0, 1.0, -0.07)

	glColor4f(0.0, 0.0, 0.0, 1.0) # blade
	glVertex3f(0.0, 1.0, 0.0)		
	glVertex3f(0.0, 1.0, -1.0)
	glColor4f(0.0, 0.0, 0.0, 0.5)	
	glVertex3f(-0.07, 1.0, -1.0)
	glVertex3f(-0.07, 1.0, -0.2)
	glColor4f(1.0, 1.0, 0.0, 1.0) # tip	
	glVertex3f(0.0, 1.0, -1.0)
	glVertex3f(0.0, 1.0, -1.04)
	glColor4f(1.0, 1.0, 0.0, 0.0)	
	glVertex3f(-0.07, 1.0, -1.04)
	glVertex3f(-0.07, 1.0, -1.0)


world = ode.World()
world.setGravity((0, -9.81, 0)) # Simulate Earth's gravity

class HeliLocalControl(object):
	
	# Control states (-1.0 to 1.0)
	collective = 0.7 # Main rotor thrust
	cyclic_pitch = 0.0 # Main rotor angle
	cyclic_roll = 0.0
	pedal = 0.0 # Tail-rotor thrust

	cyclic_mouse_sensitivity = 0.0015
	cyclic_keyboard_speed = 0.01
	collective_wheel_sensitivity = 0.02
	collective_keyboard_speed = 0.25
	pedal_keyboard_speed = 0.05
	pedal_centering = 0.75


	key_map = {
		'left_pedal': key.A,
		'right_pedal': key.D,
		'collective_down': key.W,
		'collective_up': key.S,
		'cyclic_left': key.LEFT,
		'cyclic_right': key.RIGHT,
		'cyclic_forward': key.UP,
		'cyclic_back': key.DOWN,
		'pause': key.P,
	}

	def __init__(self, window):
		self.configure_keys()
		window.push_handlers(self)
	
	def configure_keys(self):
		"""Configure key map"""
		self.active_functions = set()
		self.key2func = {}
		for funcname, key in self.key_map.items():
			self.key2func[key] = getattr(self, funcname)

	@staticmethod
	def clamp(value, minval=-1.0, maxval=1.0):
		return min(max(value, minval), maxval)

	def on_mouse_motion(self, x, y, dx, dy):
		# Mouse controls cyclic
		self.cyclic_pitch = self.clamp(self.cyclic_pitch - dy * self.cyclic_mouse_sensitivity) 
		self.cyclic_roll = self.clamp(self.cyclic_roll - dx * self.cyclic_mouse_sensitivity) 
	
	def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
		self.collective = self.clamp(
			self.collective - scroll_y * self.collective_wheel_sensitivity, minval=0)
	
	def on_key_press(self, symbol, modifier):
		if symbol in self.key2func:
			self.active_functions.add(self.key2func[symbol])
			return True
	
	def on_key_release(self, symbol, modifier):
		try:
			self.active_functions.remove(self.key2func[symbol])
			return True
		except KeyError:
			pass
	
	def update(self, dt):
		if abs(self.pedal) < self.pedal_keyboard_speed / 2:
			self.pedal = 0
		'''
		if abs(self.cyclic_pitch) < self.cyclic_keyboard_speed / 2:
			self.cyclic_pitch = 0
		if abs(self.cyclic_roll) < self.cyclic_keyboard_speed / 2:
			self.cyclic_roll = 0
		'''

		for func in self.active_functions:
			func(dt)
		self.pedal *= self.pedal_centering
		#self.cyclic_pitch *= self.pedal_centering
		#self.cyclic_roll *= self.pedal_centering

	# Control functions #

	def left_pedal(self, dt):
		self.pedal = self.clamp(self.pedal + self.pedal_keyboard_speed)
	
	def right_pedal(self, dt):
		self.pedal = self.clamp(self.pedal - self.pedal_keyboard_speed)
	
	def collective_down(self, dt):
		self.collective = self.clamp(
			self.collective - self.collective_keyboard_speed * dt, minval=0)
	
	def collective_up(self, dt):
		self.collective = self.clamp(
			self.collective + self.collective_keyboard_speed * dt, minval=0)
	
	def cyclic_left(self, dt):
		self.cyclic_roll = self.clamp(self.cyclic_roll + self.cyclic_keyboard_speed)
	
	def cyclic_right(self, dt):
		self.cyclic_roll = self.clamp(self.cyclic_roll - self.cyclic_keyboard_speed)
	
	def cyclic_forward(self, dt):
		self.cyclic_pitch = self.clamp(self.cyclic_pitch - self.cyclic_keyboard_speed)

	def cyclic_back(self, dt):
		self.cyclic_pitch = self.clamp(self.cyclic_pitch + self.cyclic_keyboard_speed)
	
	def pause(self, dt):
		global PAUSED
		PAUSED = not PAUSED


class HeliBody(ode.Body):

	mass = 7500 # kg
	rotor_max_thrust = 3500000.0
	cyclic_thrust = rotor_max_thrust / 10.0
	rotor_centering_force = cyclic_thrust / 10.0
	tail_rotor_thrust = rotor_max_thrust / 10.0
	max_vert_speed = 10.0

	# Physical dimensions
	length = 18.0
	width = 4.0
	height = 4.5
	main_rotor_height = height * 0.75
	main_rotor_radius = 5.1
	tail_rotor_pos = (0.0, 0.0, length - 4.0)

	# Aerodynamic drag constants
	frontal_Cd = 0.4
	frontal_area = width * height * 0.75
	zdrag_factor = 10.0 * frontal_Cd * frontal_area  / 2.0
	side_Cd = 0.9
	side_area = height * length * 0.55
	xdrag_factor = 10.0 * side_Cd * side_area / 2.0
	xdrag_pos = (0.0, 0.0, length / 50.0)
	top_Cd = 0.75
	top_area = width * length * 0.4
	ydrag_factor = 10.0 * top_Cd * top_area / 2.0
	ydrag_pos = (0.0, 0.0, length / 50.0)

	def __new__(cls, control):
		return ode.Body.__new__(cls, world)
	
	def __init__(self, control):
		ode.Body.__init__(self, world)
		self.control = control
		mass = ode.Mass()
		mass.setSphereTotal(self.mass, self.width / 2)
		self.setMass(mass)

	def update(self, dt):
		self.control.update(dt)
		avx, avy, avz = self.getAngularVel()
		x, y, z = self.getPosition()

		# Damp angular velocity, this simulates various angular drag forces
		# and gyroscopic forces and keeps things stable
		self.setAngularVel((avx * 0.90, avy * 0.90, avz * 0.90))

		#ravx, ravy, ravz = self.getRelAngularVel()
		rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = self.getRotation()
		rvx, rvy, rvz = self.getRelLinearVel()

		# Calculate air density at this altitude
		air_density = AIR_DENSITY * exp((-9.81 * 0.0289644 * y) / (8.31432 * AIR_TEMP))

		# Calculate aerodynamic drag and "weathervane" forces
		zdrag = (rvz)**2 * self.zdrag_factor * air_density
		xdrag = (rvx)**2 * self.xdrag_factor * air_density
		ydrag = (rvy)**2 * self.ydrag_factor * air_density
		if rvx > 0.0:
			xdrag = -xdrag
		if rvy > 0.0:
			ydrag = -ydrag
		if rvz > 0.0:
			zdrag = -zdrag
		self.addRelForceAtRelPos((xdrag, 0.0, 0.0), self.xdrag_pos)
		self.addRelForceAtRelPos((0.0, ydrag, 0.0), self.ydrag_pos)
		self.addRelForce((0.0, 0.0, zdrag))

		# rotor lift is enhanced by parallel relative wind
		# The lift goes negative at low collective
		# Which can be useful in aerobatic maneuvers
		sf = abs(rvz) / 150.0
		lift_factor = (3*sf**2 + sf**4 - 0.05) / (6 * sf**3 + 3*sf**2 + 1.0) + 1.05
		vert_collective_adjust = rvy / self.max_vert_speed * -0.15
		lift = ((self.control.collective+vert_collective_adjust - 0.1) * self.rotor_max_thrust * lift_factor * air_density**2 + zdrag) / 4.0

		# calculate the rotor lift at four positions, modified by the 
		# collective and cyclic control state and angular momentum drag
		#lift = self.rotor_max_thrust * self.control.collective / 4.0
		# Upright
		roll_centering = 0 #rxy**3 * self.rotor_centering_force
		pitch_centering = 0 #rzy**3 * self.rotor_centering_force
		if ryy < 0.0:
			# inverted
			roll_centering = -roll_centering
			pitch_centering = -pitch_centering
		pitch = self.control.cyclic_pitch*abs(self.control.cyclic_pitch)
		roll = self.control.cyclic_roll*abs(self.control.cyclic_roll)
		self.addRelForceAtRelPos( # Front blade force
			(0.0, lift - pitch * self.cyclic_thrust - pitch_centering, -lift*0.01),
			(0.0, self.main_rotor_height, self.main_rotor_radius))
		self.addRelForceAtRelPos( # Rear blade force
			(0.0, lift + pitch * self.cyclic_thrust + pitch_centering, lift*0.01),
			(0.0, self.main_rotor_height, -self.main_rotor_radius))
		self.addRelForceAtRelPos( # Port blade force
			(lift*0.01, lift - roll * self.cyclic_thrust + roll_centering, 0),
			(-self.main_rotor_radius, self.main_rotor_height, 0.0))
		self.addRelForceAtRelPos( # Starboard blade force
			(-lift*0.01, lift + roll * self.cyclic_thrust - roll_centering, 0),
			(self.main_rotor_radius, self.main_rotor_height, 0.0))

		self.addRelForceAtRelPos( # Tail rotor
			(self.control.pedal * self.tail_rotor_thrust, 0.0, 0.0),
			self.tail_rotor_pos)
		#self.addRelTorque((0.0, self.control.pedal * self.pedal_torque, 0.0))

	def mv_matrix(self):
		"""Return the 4x4 opengl model view matrix for this body as a ctypes array
		of 16 floats
		"""
		rx1, ry1, rz1, rx2, ry2, rz2, rx3, ry3, rz3 = self.getRotation()
		px, py, pz = self.getPosition()
		return (GLfloat * 16)(
			rx1, rx2, rx3, 0.0,
			ry1, ry2, ry3, 0.0,
			rz1, rz2, rz3, 0.0,
			px, py, pz, 1.0)
	
	def getRelLinearVel(self):
		"""Return the linear velocity as a 3-tuple relative to the body's frame of reference"""
		lvx, lvy, lvz = self.getLinearVel()
		rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = self.getRotation()
		return (
			lvx*rxx + lvy*rxy + lvz*rxz, 
			lvx*ryx + lvy*ryy + lvz*ryz, 
			lvx*rzx + lvy*rzy + lvz*rzz)
	
	def getRelAngularVel(self):
		"""Return the angular velocity as a 3-tuple relative to the body's frame of reference"""
		avx, avy, avz = self.getAngularVel()
		rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = self.getRotation()
		return (
			avx*rxx + avy*ryx + avz*rzx, 
			avx*rxy + avy*ryy + avz*rzy, 
			avx*rxz + avy*ryz + avz*rzz)
	
	def setRelAngularVel(self, (avx, avy, avz)):
		rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = self.getRotation()
		self.setAngularVel((
			avx*rxx + avy*ryx + avz*rzx, 
			avx*rxy + avy*ryy + avz*rzy, 
			avx*rxz + avy*ryz + avz*rzz))


class ModelView(object):
	"""Generic 3d model viewer"""
	
	_models = {} # Class var

	def __init__(self, model_name):
		self.model = data.load_model(model_name)
	
	
	def draw(self, body):
		"""Draw the model with the position and orientation of the specified body"""
		glDisable(GL_COLOR_MATERIAL)
		glPushMatrix()
		glMultMatrixf(body.mv_matrix())
		self.model.draw()
		glEnable(GL_COLOR_MATERIAL)
		glTranslatef(0.0, 1.3, 0.0)
		draw_rotor(body.main_rotor_radius, 0.15)
		glTranslatef(-0.1, -0.4, 6.0)
		glRotatef(90.0, 0.0, 0.0, 1.0)
		draw_rotor(1.0)

		glPopMatrix()


class Camera(object):

	def __init__(self, target, offset_x, offset_y, offset_z):
		self.target = target
		self.offset_x = offset_x
		self.offset_y = offset_y
		self.offset_z = offset_z
		self.bearing = 0.0
		self.pitch = 0.0
	
	def update(self):
		rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = self.target.getRotation()
		vx, vy, vz = self.target.getLinearVel()
		vm = sqrt(vx**2 + vz**2)
		if vm > 1.0:
			vx /= vm
			vz /= vm
		self.bearing = degrees(atan2(rzx - vx/3.0, rzz - vz/3.0))
	
	def apply(self):
		glLoadIdentity()
		glTranslatef(self.offset_x , self.offset_y, self.offset_z)
		glRotatef(-self.bearing, 0, 1, 0)
		px, py, pz = self.target.getPosition()
		glTranslatef(-px, -py, -pz)


font = pyglet.font.load(None, 12)

def begin_2d():
	glPushAttrib(GL_ENABLE_BIT)
	glDisable(GL_DEPTH_TEST)
	glDisable(GL_LIGHTING)
	glMatrixMode(GL_PROJECTION)
	glPushMatrix()
	glLoadIdentity()
	glOrtho(0, win.width, 0, win.height, -1, 1)
	glMatrixMode(GL_MODELVIEW)
	glPushMatrix()
	glLoadIdentity()

def end_2d():
	glPopMatrix()
	glMatrixMode(GL_PROJECTION)
	glPopMatrix()
	glMatrixMode(GL_MODELVIEW)
	glPopAttrib()

class Dashboard(object):

	interval = 0.33

	def __init__(self, body):
		self.body = body
		self.next_update = 0.0
		self.time = 0.0
		self.text = []
	
	def update(self, dt):
		self.time += dt
		if self.time > self.next_update:
			px, py, pz = self.body.getPosition()
			rvx, rvy, rvz = self.body.getRelLinearVel()
			vx, vy, vz = self.body.getLinearVel()
			self.text = [
				pyglet.font.Text(font, 'alt: %d' % py,
					x=5, y=12, halign=pyglet.font.Text.LEFT, valign=pyglet.font.Text.CENTER),
				pyglet.font.Text(font, 'spd: %d' % -rvz,
					x=100, y=12, halign=pyglet.font.Text.LEFT, valign=pyglet.font.Text.CENTER),
				pyglet.font.Text(font, 'vspd: %d' % vy,
					x=200, y=12, halign=pyglet.font.Text.LEFT, valign=pyglet.font.Text.CENTER),
			]
			self.height = max(t.height for t in self.text) 
			self.next_update = self.time + self.interval
	
	def draw(self):
		begin_2d()
		glColor4f(0.0, 0.0, 0.0, 0.7)
		glBegin(GL_QUADS)
		glVertex3f(0.0, 0.0, 0.0)
		glVertex3f(win.width, 0.0, 0.0)
		glVertex3f(win.width, self.height * 1.5, 0.0)
		glVertex3f(0.0, self.height * 1.5, 0.0)
		glEnd()
		for t in self.text:
			t.draw()
		end_2d()


def main():
	global world, win, PAUSED
	config_templates = [
		Config(sample_buffers=1, samples=4, stencil_size=1,
						depth_size=24, double_buffer=True),
		Config(sample_buffers=1, samples=4, stencil_size=1,
						depth_size=16, double_buffer=True),
		Config(stencil_size=1, depth_size=24, double_buffer=True),
		Config(stencil_size=1, depth_size=16, double_buffer=True),
	]

	screen = pyglet.window.get_platform().get_default_display().get_default_screen()

	for template in config_templates:
		try:
			config = screen.get_best_config(template)
		except pyglet.window.NoSuchConfigException:
			continue
		else:
			break
	else:
		raise window.NoSuchConfigException()

	win = window.Window(visible=False, config=config, fullscreen=True)
	win.on_resize=resize
	win.set_exclusive_mouse()
	heli_body = HeliBody(HeliLocalControl(win))
	heli_body.setPosition((0, 1000, 0))
	heli_view = ModelView('heli.obj')
	camera = Camera(heli_body, 0.0, -6.0, -25.0)
	ground = terrain.FractalTerrainMesh(iterations=7, deviation=500.0, smooth=3.0,
		display_width=25000, snow_line=600.0, tree_line=550.0, water_line=-80)
	ground.compile()
	dash = Dashboard(heli_body)
	skybox = SkyBox.fromDir(data.texture_dir + '/plainsky', 'plainsky')

	init()

	win.set_visible()
	clock = pyglet.clock.Clock(30)
	clock.tick() # Prime the clock so our first time delta isn't 0

	while not win.has_exit:
		dt=clock.tick()
		win.dispatch_events()
		if not PAUSED:
			heli_body.update(dt)
			dash.update(dt)
			world.quickStep(dt)
		
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT)
		#glTranslatef(0.0, -3.0, -12.0)
		#heli_view.model.draw()
		camera.update()
		skybox.draw(-camera.pitch, -camera.bearing)
		camera.apply()

		glPushMatrix()
		glRotatef(-90, 1, 0, 0)
		ground.draw_composed()
		glPopMatrix()
		heli_view.draw(heli_body)
		dash.draw()

		win.flip()

	
	print "fps:  %d" % clock.get_fps()

if __name__ == '__main__': main()
